#include <absl/flags/flag.h>
#include <absl/flags/parse.h>

#include <algorithm>
#include <fstream>
#include <iostream>
#include <sstream>
#include <vector>

/**
 * @brief 载入真值、测试结果
 *
 * @param filename
 * @return std::vector<std::vector<int>>
 */
std::vector<std::vector<int>> load_from_file(const std::string& filename);

/**
 * @brief 用于保存真值、测试结果
 *
 * @param vec
 * @param filename
 */
void save_to_file(const std::vector<std::vector<int>>& vec, const std::string& filename);

/**
 * @brief 查找vector_b中是否包含vector_a中的元素
 *
 * @param vector_a
 * @param vector_b
 * @param rank
 * @param out_file
 * @return true
 * @return false
 */
bool find_common_elements(const std::vector<int>& vector_a,
                          const std::vector<int>& vector_b,
                          int rank,
                          std::ofstream& out_file);

ABSL_FLAG(int, rank, 5, "Rank");
ABSL_FLAG(std::string, ground_truth, "../data/rank_5_ground_truth_test.txt", "ground_truth path");
ABSL_FLAG(std::string, result, "../data/rank_5_result_test.txt", "reid_result path");
ABSL_FLAG(std::string, result_log, "../data/rank_5_evaluation_test.txt", "result_log path");
ABSL_FLAG(int, dataset_num, 200, "数据集数量");
ABSL_FLAG(int, test_num, 20, "测试集数量");

int main(int argc, char** argv) {
  // 1.0 载入配置
  absl::ParseCommandLine(argc, argv);
  std::cout << "usage: ./rank_evaluation "
            << "--rank=5 "
            << "--ground_truth=../data/rank_5_ground_truth_test.txt "
            << "--result=../data/rank_5_result_test.txt "
            << "--result_log=../data/rank_5_evaluation_test.txt "
            << "--dataset_num=200 "
            << "--test_num=20 "
            << std::endl;
  int N_rank = absl::GetFlag(FLAGS_rank);
  int N_dataset = absl::GetFlag(FLAGS_dataset_num);  // Gallery图片数量
  int N_test = absl::GetFlag(FLAGS_test_num);        // Query图片数量
  std::string path_gt = absl::GetFlag(FLAGS_ground_truth);
  std::string path_rr = absl::GetFlag(FLAGS_result);
  std::string path_evaluation = absl::GetFlag(FLAGS_result_log);

  // 2.0 载入真值与测试值
  std::vector<std::vector<int>> ground_truth;
  std::vector<std::vector<int>> reid_result;
  ground_truth = load_from_file(path_gt);
  reid_result = load_from_file(path_rr);
  // 2.1 丢弃测试值超过Rank的数据
  for (int i = 0; i < reid_result.size(); i++)
    while (reid_result[i].size() > N_rank + 1)
      reid_result[i].pop_back();
  // 2.2 测试数据不符合要求时无法计算
  if (reid_result.size() != N_test)
    std::cout << "reid_result size error: " << reid_result.size() << std::endl;

  // 3.0 结果计算
  // 3.1 初始化结果计算
  std::vector<bool> reid_success(N_test, false);
  int success_count = 0;
  std::ofstream evaluation_file(path_evaluation);
  if (!evaluation_file.is_open()) {
    std::cerr << "Error opening file for writing: \n"
              << path_evaluation << std::endl;
    return 0;
  }
  evaluation_file << "input ground truth file: " << path_gt << std::endl;
  evaluation_file << "input result file: " << path_rr << std::endl;
  evaluation_file << "N_dataset: " << N_dataset << std::endl;
  evaluation_file << "N_test: " << N_test << std::endl;
  evaluation_file << "N_rank: " << N_rank << std::endl
                  << std::endl;

  // 3.2 评估每个目标
  for (int i = 0; i < N_test; i++) {
    std::cout << "target " << i << std::endl;
    evaluation_file << "target " << i << std::endl;
    // 3.2.1 Rank查找
    if (find_common_elements(reid_result[i], ground_truth[i], N_rank, evaluation_file)) {
      reid_success[i] = true;
      success_count++;
    } else {
      std::cout << "Element not found in ground_truth" << std::endl;
      evaluation_file << "Element not found in ground_truth" << std::endl;
    }
  }

  // 3.3 结果计算
  double score = (success_count * 1.0) / (N_test * 1.0);
  std::cout << "score: " << score << std::endl;
  evaluation_file << "score: " << score << " correct result: " << success_count << std::endl;
  evaluation_file.close();

  return 0;
}

std::vector<std::vector<int>> load_from_file(const std::string& filename) {
  std::ifstream infile(filename);
  if (!infile.is_open()) {
    std::cerr << "Error opening file for reading" << std::endl;
    return {};
  }
  std::vector<std::vector<int>> vec;
  std::string line;
  // 按行读取文件
  while (std::getline(infile, line)) {
    std::stringstream ss(line);
    std::vector<int> row;
    int num;
    // 读取每行的整数并放入row中
    while (ss >> num) {
      row.push_back(num);
    }
    vec.push_back(row);
  }
  infile.close();
  return vec;
}

void save_to_file(const std::vector<std::vector<int>>& vec, const std::string& filename) {
  std::ofstream outfile(filename);
  if (!outfile.is_open()) {
    std::cerr << "Error opening file for writing" << std::endl;
    return;
  }
  // 将每行的数据写入文件
  for (const auto& row : vec) {
    for (const auto& elem : row) {
      outfile << elem << " ";  // 用空格分隔每个元素
    }
    outfile << "\n";  // 每一行结束后换行
  }
  outfile.close();
}

bool find_common_elements(const std::vector<int>& vector_a,
                          const std::vector<int>& vector_b,
                          int rank,
                          std::ofstream& out_file) {
  // 遍历 vector_a 中的每一个元素
  bool if_found = false;
  for (size_t i = 1; i < vector_a.size() && i < rank + 1; ++i) {
    // 查找 vector_a[i] 是否在 vector_b 中
    auto it = std::find(vector_b.begin(), vector_b.end(), vector_a[i]);
    // 如果找到了，输出 vector_a 的位置以及在 vector_b 中的位置
    if (it != vector_b.end()) {
      size_t pos_in_b = std::distance(vector_b.begin(), it);  // 计算位置
      std::cout << "Element " << vector_a[i] << " found at position " << i
                << " in reid_result and position " << pos_in_b << " in ground_truth." << std::endl;
      out_file << "Element " << vector_a[i] << " found at position " << i
               << " in reid_result and position " << pos_in_b << " in ground_truth." << std::endl;
      if_found = true;
    }
  }
  return if_found;
}